"""最小二乘法"""
import numpy as np
import datetime
import tushare as ts
from chinese_calendar import is_holiday
from numpy_ext import rolling_apply
# from stockSelect10_soup_phone_1  import getstock
# import pandas as pd
# import time

"""
def masscenter(price, nQty):
    return np.sum(price * nQty) / np.sum(nQty)


df = pd.DataFrame( [['02:59:47.000282', 87.60, 739],
                    ['03:00:01.042391', 87.51, 10],
                    ['03:00:01.630182', 87.51, 10],
                    ['03:00:01.635150', 88.00, 792],
                    ['03:00:01.914104', 88.00, 10]], 
                   columns=['stamp', 'price','nQty'])
df['stamp'] = pd.to_datetime(df['stamp'], format='%H:%M:%S.%f')
df.set_index('stamp', inplace=True, drop=True)

window = 2
df['y'] = rolling_apply(masscenter, window, df.price.values, df.nQty.values)
print(df)
"""
# import matplotlib.pyplot as plt

ts.set_token("6667cd4a2326f2f937062a0f4fb59aea5c56d13b1f6f26225f115fe9")
pro = ts.pro_api()

# def fun2ploy(x,n):  #得到范德蒙德（Vandermonde）行列式,#https://zhuanlan.zhihu.com/p/38128785/
#     '''
#     数据转化为[x^0,x^1,x^2,...x^n]
#     首列变1
#     '''
#     lens = len(x)
#     X = np.ones([1,lens])
#     for i in range(1,n):
#         X = np.vstack((X,np.power(x,i)))#按行堆叠
#     #print('Vandermonde行列式:',X)
#     return X  

# def leastseq_byploy(x,y,ploy_dim=2):  #https://blog.csdn.net/dz4543/article/details/85224391
#     '''
#     最小二乘求解
#     '''
#     X = fun2ploy(x,ploy_dim)
#     #直接求解
#     Xt = X.transpose();#转置变成列向量 X=[[x^0][x^1][x^2][x^3]]
#     XXt=X.dot(Xt);#矩阵乘,X*X^T
#     XXtInv = np.linalg.inv(XXt)#求逆 (X*X^T)^-1
#     XXtInvX = XXtInv.dot(X) #(X*X^T)^-1*X
#     #获得系数列向量
#     coef = XXtInvX.dot(y.T)
#     # y_est = Xt.dot(coef)
#     # return y_est,coef
#     return coef


# def fit_func(p, x):  # 如 p=numpy.poly1d([1,2,3])  生成  $1x^2+2x^1+3x^0$*,coef 傅立叶展开
#     p=p[::-1]
#     f = np.poly1d(p) 
#     return f(x)

def n_days_before(st_day,n):
    """
    st_day:开始日期
    :type st_day: datetime.date | datetime.datetime
    """
    while(n>0):
        if not is_holiday(st_day):
            n=n-1
        st_day=st_day-datetime.timedelta(1)
    return st_day



def leastseq_merge(xi,yi,ploy_dim=3):  #leastseq_byploy+fun2ploy+fit_func 合并给zong_he.rolling_apply
    '''
    最小二乘求解
    '''
    x=xi[:-1] #去掉最后一个，作为当天数据，以window-1个数据作预测系数计算，xi最后一个数作预测参数
    y=yi[:-1]
    X = np.ones([1,len(x)])
    for i in range(1,ploy_dim):
        X = np.vstack((X,np.power(x,i)))
    #直接求解
    Xt = X.transpose();#转置变成列向量 X=[[x^0][x^1][x^2][x^3]]
    XXt=X.dot(Xt);#矩阵乘,X*X^T
    XXtInv = np.linalg.inv(XXt)#求逆 (X*X^T)^-1
    XXtInvX = XXtInv.dot(X) #(X*X^T)^-1 *X
    #获得系数列向量
    coef = XXtInvX.dot(y.T) 
    # y_est = Xt.dot(coef)
    # return y_est,coef
    coef=coef[::-1] #coef倒序是因为coefficients是从低次幂到高次幂，poly1d是从高到低
    f = np.poly1d(coef) #np.poly1d([1,2,3])(x)=x^2+x^1+x^0
    f_predict=f(xi[-1])
     #h_coef,LOPEN
    return f_predict #y=coef[0]+coef[1]*x

def zong_he(code=799,days=60,ploy_dim=2,window=2):#从N天前到上一天的数据，DATA[0]是最近的
    # end_date_str=time.strftime("%Y%m%d") 
    # end_datetime = datetime.datetime.strptime(end_date_str, '%Y%m%d') 
    end_datetime=datetime.date.today()
    start_datetime=n_days_before(end_datetime,days) 
    end_date_str=end_datetime.strftime("%Y%m%d") 
    start_date_str=start_datetime.strftime("%Y%m%d") 
    code6=str(code).zfill(6)
    if str(code).zfill(6)[:2]!="60":
            code_str='{}.SZ'.format(code6)
    else:
        code_str='{}.SH'.format(code6)
    df = pro.daily(ts_code=code_str, start_date=start_date_str,end_date=end_date_str) 
    df.set_index('trade_date',inplace=True) 
    df.sort_index(inplace=True)
    #获N天数据，由远到近
    HIGH=df["high"].values
    LOW=df["low"].values
    OPEN=df["open"].values
    # CLOSE=df["close"].values
    days=days
    code=code
    today_high_limit=df["close"].shift(1).values*1.1#涨停价
    today_low_limit=df["close"].shift(1).values*0.9#跌停价
    h_predict=rolling_apply(leastseq_merge,window+1,OPEN,HIGH,ploy_dim=ploy_dim) #window+1,是因为已知当日开盘价，预测当日最高价，系数计算用前window个日线的
    l_predict=rolling_apply(leastseq_merge,window+1,OPEN,LOW,ploy_dim=ploy_dim)
    # print('预测最高价:',np.round(fit_func(h_coef,OPEN),2))
    # print('预测最低价:',np.round(fit_func(l_coef,OPEN),2))
    # err_h=np.sum((fit_func(h_coef,OPEN)-HIGH)**2)/len(OPEN)
    # err_l=np.sum((fit_func(l_coef,OPEN)-LOW)**2)/len(OPEN)
    
    # print("high标准差",np.sqrt(err_h))
    # print("low标准差",np.sqrt(err_l))
    #https://baike.baidu.com/item/黑塞矩阵/2248782?fr=aladdin
    h_predict=[today_high_limit[i] if h_predict[i] > today_high_limit[i]  else h_predict[i]  for i in range(len(h_predict))]
    l_predict=[today_low_limit[i] if l_predict[i] < today_low_limit[i]  else l_predict[i]  for i in range(len(l_predict))]     
    return [HIGH,h_predict,today_high_limit,LOW,l_predict,today_low_limit]

# 最小二乘拟合
#https://www.cnblogs.com/lc1217/p/6514734.html
if __name__=='__main__':
    code=799
    days=20
    res=zong_he(code=code,days=days,ploy_dim=2,window=5)
    from _utinity import draw_,get_data
    df=get_data(code,N=days)
    draw_(df,*res)
    
    
